#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
EDR Detection and Evasion Scanner
EDR检测和规避扫描器

该模块实现：
1. 自动检测目标环境中的EDR产品
2. 根据检测结果选择对应的规避策略
3. 动态更新EDR特征库
4. 提供规避建议和实施方案
"""

from __future__ import print_function
from __future__ import unicode_literals
from __future__ import division
from __future__ import absolute_import

import os
import re
import sys
import json
import glob
import time
import subprocess
from collections import defaultdict, OrderedDict

# Python 2/3 兼容性
if sys.version_info[0] == 2:
    string_types = basestring
    text_type = unicode
else:
    string_types = str
    text_type = str


class EDRScanner(object):
    """EDR检测扫描器"""
    
    def __init__(self, signatures_path=None):
        self.signatures_path = signatures_path or self._get_default_signatures_path()
        self.signatures = {}
        self.scan_results = {}
        self.detected_edrs = []
        self.confidence_threshold = 0.7
        self.load_signatures()
    
    def _get_default_signatures_path(self):
        """获取默认特征库路径"""
        return os.path.join(
            os.path.dirname(__file__), '..', '..', 
            'config', 'edr_signatures.json'
        )
    
    def load_signatures(self):
        """加载EDR特征库"""
        try:
            if os.path.exists(self.signatures_path):
                with open(self.signatures_path, 'r', encoding='utf-8') as f:
                    self.signatures = json.load(f)
                print(f"[+] 已加载EDR特征库: {self.signatures_path}")
                print(f"[+] 特征库版本: {self.signatures.get('version', 'Unknown')}")
                print(f"[+] 包含 {len(self.signatures.get('edr_products', {}))} 个EDR产品特征")
            else:
                print(f"[!] 特征库文件不存在: {self.signatures_path}")
                self._create_minimal_signatures()
        except Exception as e:
            print(f"[!] 加载特征库失败: {e}")
            self._create_minimal_signatures()
    
    def _create_minimal_signatures(self):
        """创建最小特征库"""
        self.signatures = {
            "version": "1.0.0",
            "edr_products": {},
            "generic_detection_patterns": {
                "common_edr_indicators": [],
                "suspicious_processes": [],
                "monitoring_tools": []
            }
        }
    
    def scan_processes(self):
        """扫描进程以检测EDR"""
        print("[*] 扫描进程中的EDR特征...")
        process_results = defaultdict(list)
        
        try:
            # 获取所有进程
            ps_output = subprocess.check_output(['ps', 'aux'], universal_newlines=True)
            processes = ps_output.strip().split('\n')[1:]  # 跳过标题行
            
            # 检查每个EDR产品的进程特征
            for edr_name, edr_info in self.signatures.get('edr_products', {}).items():
                detection_methods = edr_info.get('detection_methods', {})
                target_processes = detection_methods.get('processes', [])
                
                for process_line in processes:
                    process_parts = process_line.split()
                    if len(process_parts) >= 11:
                        command = ' '.join(process_parts[10:])
                        
                        for target_process in target_processes:
                            if target_process.lower() in command.lower():
                                process_results[edr_name].append({
                                    'type': 'process',
                                    'indicator': target_process,
                                    'found': command,
                                    'confidence': 0.8
                                })
                                
        except subprocess.CalledProcessError as e:
            print(f"[!] 进程扫描失败: {e}")
        except Exception as e:
            print(f"[!] 进程扫描异常: {e}")
        
        return dict(process_results)
    
    def scan_files(self):
        """扫描文件系统以检测EDR"""
        print("[*] 扫描文件系统中的EDR特征...")
        file_results = defaultdict(list)
        
        try:
            # 检查每个EDR产品的文件特征
            for edr_name, edr_info in self.signatures.get('edr_products', {}).items():
                detection_methods = edr_info.get('detection_methods', {})
                target_files = detection_methods.get('files', [])
                
                for target_file in target_files:
                    # 支持通配符匹配
                    if '*' in target_file or '?' in target_file:
                        matching_files = glob.glob(target_file)
                        for matching_file in matching_files:
                            if os.path.exists(matching_file):
                                file_results[edr_name].append({
                                    'type': 'file',
                                    'indicator': target_file,
                                    'found': matching_file,
                                    'confidence': 0.9
                                })
                    else:
                        if os.path.exists(target_file):
                            file_results[edr_name].append({
                                'type': 'file',
                                'indicator': target_file,
                                'found': target_file,
                                'confidence': 0.9
                            })
                            
        except Exception as e:
            print(f"[!] 文件扫描异常: {e}")
        
        return dict(file_results)
    
    def scan_services(self):
        """扫描系统服务以检测EDR"""
        print("[*] 扫描系统服务中的EDR特征...")
        service_results = defaultdict(list)
        
        try:
            # 获取所有服务状态
            systemctl_output = subprocess.check_output(
                ['systemctl', 'list-units', '--type=service', '--all'], 
                universal_newlines=True
            )
            services = systemctl_output.strip().split('\n')
            
            # 检查每个EDR产品的服务特征
            for edr_name, edr_info in self.signatures.get('edr_products', {}).items():
                detection_methods = edr_info.get('detection_methods', {})
                target_services = detection_methods.get('services', [])
                
                for service_line in services:
                    service_parts = service_line.split()
                    if len(service_parts) >= 1:
                        service_name = service_parts[0]
                        
                        for target_service in target_services:
                            if target_service.lower() in service_name.lower():
                                service_results[edr_name].append({
                                    'type': 'service',
                                    'indicator': target_service,
                                    'found': service_name,
                                    'confidence': 0.85
                                })
                                
        except subprocess.CalledProcessError as e:
            print(f"[!] 服务扫描失败: {e}")
        except Exception as e:
            print(f"[!] 服务扫描异常: {e}")
        
        return dict(service_results)
    
    def scan_kernel_modules(self):
        """扫描内核模块以检测EDR"""
        print("[*] 扫描内核模块中的EDR特征...")
        module_results = defaultdict(list)
        
        try:
            # 获取已加载的内核模块
            lsmod_output = subprocess.check_output(['lsmod'], universal_newlines=True)
            modules = lsmod_output.strip().split('\n')[1:]  # 跳过标题行
            
            # 检查每个EDR产品的内核模块特征
            for edr_name, edr_info in self.signatures.get('edr_products', {}).items():
                detection_methods = edr_info.get('detection_methods', {})
                target_modules = detection_methods.get('kernel_modules', [])
                
                for module_line in modules:
                    module_parts = module_line.split()
                    if len(module_parts) >= 1:
                        module_name = module_parts[0]
                        
                        for target_module in target_modules:
                            if target_module.lower() in module_name.lower():
                                module_results[edr_name].append({
                                    'type': 'kernel_module',
                                    'indicator': target_module,
                                    'found': module_name,
                                    'confidence': 0.95
                                })
                                
        except subprocess.CalledProcessError as e:
            print(f"[!] 内核模块扫描失败: {e}")
        except Exception as e:
            print(f"[!] 内核模块扫描异常: {e}")
        
        return dict(module_results)
    
    def scan_network_indicators(self):
        """扫描网络指标以检测EDR"""
        print("[*] 扫描网络连接中的EDR特征...")
        network_results = defaultdict(list)
        
        try:
            # 获取网络连接
            netstat_output = subprocess.check_output(
                ['netstat', '-tuln'], 
                universal_newlines=True
            )
            connections = netstat_output.strip().split('\n')
            
            # 检查DNS解析记录
            try:
                resolv_conf = '/etc/resolv.conf'
                if os.path.exists(resolv_conf):
                    with open(resolv_conf, 'r') as f:
                        dns_content = f.read()
                        
                        # 检查每个EDR产品的网络特征
                        for edr_name, edr_info in self.signatures.get('edr_products', {}).items():
                            detection_methods = edr_info.get('detection_methods', {})
                            network_indicators = detection_methods.get('network_indicators', [])
                            
                            for indicator in network_indicators:
                                # 简化的域名匹配
                                domain_pattern = indicator.replace('*', '').replace('.', r'\.')
                                if domain_pattern in dns_content:
                                    network_results[edr_name].append({
                                        'type': 'network',
                                        'indicator': indicator,
                                        'found': f"DNS配置中发现: {domain_pattern}",
                                        'confidence': 0.6
                                    })
            except Exception as e:
                print(f"[!] DNS检查失败: {e}")
                
        except subprocess.CalledProcessError as e:
            print(f"[!] 网络扫描失败: {e}")
        except Exception as e:
            print(f"[!] 网络扫描异常: {e}")
        
        return dict(network_results)
    
    def scan_generic_indicators(self):
        """扫描通用EDR指标"""
        print("[*] 扫描通用EDR指标...")
        generic_results = []
        
        try:
            generic_patterns = self.signatures.get('generic_detection_patterns', {})
            
            # 检查通用EDR指标
            common_indicators = generic_patterns.get('common_edr_indicators', [])
            for indicator in common_indicators:
                if os.path.exists(indicator):
                    generic_results.append({
                        'type': 'generic_indicator',
                        'indicator': indicator,
                        'found': indicator,
                        'confidence': 0.5
                    })
            
            # 检查可疑进程
            suspicious_processes = generic_patterns.get('suspicious_processes', [])
            try:
                ps_output = subprocess.check_output(['ps', 'aux'], universal_newlines=True)
                for suspicious_proc in suspicious_processes:
                    if suspicious_proc in ps_output:
                        generic_results.append({
                            'type': 'suspicious_process',
                            'indicator': suspicious_proc,
                            'found': f"进程中发现: {suspicious_proc}",
                            'confidence': 0.4
                        })
            except:
                pass
            
            # 检查监控工具
            monitoring_tools = generic_patterns.get('monitoring_tools', [])
            for tool in monitoring_tools:
                try:
                    subprocess.check_output(['which', tool], stderr=subprocess.DEVNULL)
                    generic_results.append({
                        'type': 'monitoring_tool',
                        'indicator': tool,
                        'found': f"系统中安装了监控工具: {tool}",
                        'confidence': 0.3
                    })
                except subprocess.CalledProcessError:
                    pass
                    
        except Exception as e:
            print(f"[!] 通用指标扫描异常: {e}")
        
        return generic_results
    
    def full_scan(self):
        """执行完整的EDR扫描"""
        print("[*] 开始完整EDR扫描...")
        start_time = time.time()
        
        # 执行各种扫描
        scan_functions = [
            ('进程扫描', self.scan_processes),
            ('文件扫描', self.scan_files),
            ('服务扫描', self.scan_services),
            ('内核模块扫描', self.scan_kernel_modules),
            ('网络指标扫描', self.scan_network_indicators)
        ]
        
        all_results = {}
        
        for scan_name, scan_func in scan_functions:
            try:
                print(f"[*] 执行 {scan_name}...")
                results = scan_func()
                all_results[scan_name] = results
            except Exception as e:
                print(f"[!] {scan_name} 失败: {e}")
                all_results[scan_name] = {}
        
        # 扫描通用指标
        print("[*] 执行通用指标扫描...")
        generic_results = self.scan_generic_indicators()
        all_results['通用指标扫描'] = {'generic': generic_results}
        
        # 分析结果
        self.scan_results = all_results
        self.analyze_results()
        
        scan_time = time.time() - start_time
        print(f"[+] EDR扫描完成，耗时: {scan_time:.2f}秒")
        
        return self.scan_results
    
    def analyze_results(self):
        """分析扫描结果"""
        print("[*] 分析扫描结果...")
        
        edr_scores = defaultdict(float)
        edr_evidences = defaultdict(list)
        
        # 计算每个EDR的置信度分数
        for scan_type, scan_results in self.scan_results.items():
            if scan_type == '通用指标扫描':
                continue
                
            for edr_name, evidences in scan_results.items():
                for evidence in evidences:
                    confidence = evidence.get('confidence', 0.5)
                    edr_scores[edr_name] += confidence
                    edr_evidences[edr_name].append(evidence)
        
        # 确定检测到的EDR
        self.detected_edrs = []
        for edr_name, total_score in edr_scores.items():
            evidence_count = len(edr_evidences[edr_name])
            
            # 计算最终置信度 (考虑证据数量和质量)
            if evidence_count > 0:
                final_confidence = min(total_score / evidence_count, 1.0)
                
                if final_confidence >= self.confidence_threshold:
                    edr_info = self.signatures.get('edr_products', {}).get(edr_name, {})
                    self.detected_edrs.append({
                        'name': edr_name,
                        'display_name': edr_info.get('name', edr_name),
                        'confidence': final_confidence,
                        'evidence_count': evidence_count,
                        'evidences': edr_evidences[edr_name],
                        'evasion_strategies': edr_info.get('evasion_strategies', {})
                    })
        
        # 按置信度排序
        self.detected_edrs.sort(key=lambda x: x['confidence'], reverse=True)
        
        print(f"[+] 检测到 {len(self.detected_edrs)} 个EDR产品")
    
    def get_evasion_recommendations(self):
        """获取规避建议"""
        recommendations = []
        
        for detected_edr in self.detected_edrs:
            edr_name = detected_edr['name']
            evasion_strategies = detected_edr.get('evasion_strategies', {})
            
            edr_recommendations = {
                'edr_name': detected_edr['display_name'],
                'confidence': detected_edr['confidence'],
                'strategies': []
            }
            
            for strategy_type, strategy_info in evasion_strategies.items():
                edr_recommendations['strategies'].append({
                    'type': strategy_type,
                    'method': strategy_info.get('method', 'Unknown'),
                    'description': strategy_info.get('description', 'No description'),
                    'priority': self._get_strategy_priority(strategy_type)
                })
            
            # 按优先级排序策略
            edr_recommendations['strategies'].sort(key=lambda x: x['priority'], reverse=True)
            recommendations.append(edr_recommendations)
        
        return recommendations
    
    def _get_strategy_priority(self, strategy_type):
        """获取策略优先级"""
        priority_map = {
            'process_hiding': 10,
            'behavioral_evasion': 9,
            'file_operations': 8,
            'network_evasion': 7
        }
        return priority_map.get(strategy_type, 5)
    
    def print_scan_summary(self):
        """打印扫描摘要"""
        print("\n" + "="*60)
        print("EDR检测扫描报告")
        print("="*60)
        
        if not self.detected_edrs:
            print("[+] 未检测到已知的EDR产品")
        else:
            print(f"[!] 检测到 {len(self.detected_edrs)} 个EDR产品:")
            
            for i, edr in enumerate(self.detected_edrs, 1):
                print(f"\n{i}. {edr['display_name']}")
                print(f"   置信度: {edr['confidence']:.2%}")
                print(f"   证据数量: {edr['evidence_count']}")
                
                # 显示主要证据
                evidences = edr['evidences'][:3]  # 只显示前3个证据
                for evidence in evidences:
                    print(f"   - {evidence['type']}: {evidence['found']}")
                
                if len(edr['evidences']) > 3:
                    print(f"   - ... 还有 {len(edr['evidences']) - 3} 个证据")
        
        # 显示规避建议
        recommendations = self.get_evasion_recommendations()
        if recommendations:
            print(f"\n规避建议:")
            for rec in recommendations:
                print(f"\n针对 {rec['edr_name']}:")
                for strategy in rec['strategies'][:2]:  # 只显示前2个策略
                    print(f"  - {strategy['type']}: {strategy['description']}")
        
        print("="*60)
    
    def export_results(self, output_path=None):
        """导出扫描结果"""
        if not output_path:
            output_path = os.path.join(
                os.path.dirname(__file__), '..', '..', 
                'logs', f'edr_scan_results_{int(time.time())}.json'
            )
        
        export_data = {
            'scan_timestamp': time.time(),
            'detected_edrs': self.detected_edrs,
            'scan_results': self.scan_results,
            'recommendations': self.get_evasion_recommendations()
        }
        
        try:
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(export_data, f, indent=2, ensure_ascii=False)
            
            print(f"[+] 扫描结果已导出: {output_path}")
            return output_path
            
        except Exception as e:
            print(f"[!] 导出结果失败: {e}")
            return None


class EDREvader(object):
    """EDR规避器"""
    
    def __init__(self, scanner_results=None):
        self.scanner_results = scanner_results
        self.evasion_modules = {}
        self._load_evasion_modules()
    
    def _load_evasion_modules(self):
        """加载规避模块"""
        # 这里可以动态加载各种规避模块
        self.evasion_modules = {
            'process_hiding': self._process_hiding_module,
            'file_operations': self._file_operations_module,
            'network_evasion': self._network_evasion_module,
            'behavioral_evasion': self._behavioral_evasion_module
        }
    
    def _process_hiding_module(self, method, target_edr):
        """进程隐藏模块"""
        print(f"[*] 执行进程隐藏策略: {method}")
        # 实际的进程隐藏实现
        return True
    
    def _file_operations_module(self, method, target_edr):
        """文件操作模块"""
        print(f"[*] 执行文件操作策略: {method}")
        # 实际的文件操作实现
        return True
    
    def _network_evasion_module(self, method, target_edr):
        """网络规避模块"""
        print(f"[*] 执行网络规避策略: {method}")
        # 实际的网络规避实现
        return True
    
    def _behavioral_evasion_module(self, method, target_edr):
        """行为规避模块"""
        print(f"[*] 执行行为规避策略: {method}")
        # 实际的行为规避实现
        return True
    
    def apply_evasion_strategies(self, detected_edrs):
        """应用规避策略"""
        print("[*] 开始应用EDR规避策略...")
        
        for edr in detected_edrs:
            edr_name = edr['name']
            strategies = edr.get('evasion_strategies', {})
            
            print(f"[*] 针对 {edr['display_name']} 应用规避策略...")
            
            for strategy_type, strategy_info in strategies.items():
                method = strategy_info.get('method')
                
                if strategy_type in self.evasion_modules:
                    try:
                        success = self.evasion_modules[strategy_type](method, edr_name)
                        if success:
                            print(f"[+] {strategy_type} 策略应用成功")
                        else:
                            print(f"[!] {strategy_type} 策略应用失败")
                    except Exception as e:
                        print(f"[!] {strategy_type} 策略执行异常: {e}")
                else:
                    print(f"[!] 未找到 {strategy_type} 规避模块")


def main():
    """主函数 - 用于测试"""
    print("EDR检测和规避扫描器")
    print("="*40)
    
    # 创建扫描器
    scanner = EDRScanner()
    
    # 执行扫描
    results = scanner.full_scan()
    
    # 打印摘要
    scanner.print_scan_summary()
    
    # 导出结果
    scanner.export_results()
    
    # 如果检测到EDR，演示规避策略
    if scanner.detected_edrs:
        evader = EDREvader()
        evader.apply_evasion_strategies(scanner.detected_edrs)


if __name__ == "__main__":
    main()